from langchain.agents import Tool, AgentExecutor, create_react_agent
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain_core.messages import SystemMessage, HumanMessage
from typing import List, Dict, Any, Optional
import logging
import re
from config import LLM_CONFIG

# 导入工具模块
try:
    from tools.map_rendering import ImageRenderingTools
    image_rendering_tools = ImageRenderingTools()
except ImportError:
    image_rendering_tools = None

logger = logging.getLogger(__name__)

class GeographicMapAgent:
    """地理专题图制作专家"""
    
    def __init__(self):
        # 构建ChatOpenAI参数
        llm_params = {
            'api_key': LLM_CONFIG['openai_api_key'],
            'base_url': LLM_CONFIG['base_url'],
            'model': LLM_CONFIG['model_name'],
            'temperature': LLM_CONFIG['temperature']
        }
        # 只有当max_tokens不为None时才添加该参数
        if LLM_CONFIG['max_tokens'] is not None:
            llm_params['max_tokens'] = LLM_CONFIG['max_tokens']
        
        self.llm = ChatOpenAI(**llm_params)
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True
        )
        self.tools = self._create_tools()
        self.agent = self._create_agent()
        self.context_variables = {}  # 存储中间结果的变量
    
    def _create_tools(self) -> List[Tool]:
        """创建工具列表"""
        tools = [
            Tool(
                name="制作地理专题图",
                description="制作高质量的地理专题图，支持PNG和GIF格式。输入：配置参数JSON字符串",
                func=lambda x: self._execute_create_geographic_map(x)
            ),
            Tool(
                name="保存变量结果",
                description="将工具执行结果保存到变量中。输入：变量名，数据",
                func=lambda x: self._save_to_variable(x)
            ),
            Tool(
                name="获取变量结果",
                description="获取已保存的变量结果。输入：变量名",
                func=lambda x: self._get_from_variable(x)
            )
        ]
        return tools
    
    def _create_agent(self):
        """创建智能体"""
        prompt = PromptTemplate.from_template("""
你是一个专业的地理专题图制作专家，能够理解用户的地图制作需求，并智能解析配置参数来制作高质量的地理专题图。

当用户提出地图制作需求时，你需要：
1. 分析用户需求，提取地图配置参数
2. 解析绘制范围、比例尺、指北针、底图颜色、经纬网等信息
3. 如果用户没有指明某些参数，使用合理的默认值
4. 调用地图制作工具生成PNG或GIF格式的地图
5. 向用户展示最终结果

可用工具：
{tools}

工具名称：
{tool_names}

对话历史：
{chat_history}

用户输入：{input}

请按照以下格式思考和回应：
思考：我需要理解用户的地图制作需求，并解析配置参数
行动：选择要使用的工具
行动输入：工具所需的输入参数（JSON格式的配置参数）
观察：工具执行的结果
思考：基于结果继续下一步
... (重复这个过程直到完成任务)
最终答案：向用户展示最终结果

{agent_scratchpad}
""")
        
        agent = create_react_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=prompt
        )
        
        return AgentExecutor(
            agent=agent,
            tools=self.tools,
            memory=self.memory,
            verbose=True,
            max_iterations=5,
            handle_parsing_errors=True
        )
    
    def _save_to_variable(self, input_str: str) -> str:
        """保存数据到变量"""
        try:
            parts = input_str.split('|', 1)
            if len(parts) != 2:
                return "错误：输入格式应为 '变量名|数据'"
            
            var_name, data = parts
            var_name = var_name.strip()
            
            self.context_variables[var_name] = data.strip()
            return f"数据已保存到变量 '{var_name}'"
        except Exception as e:
            return f"保存变量失败: {str(e)}"
    
    def _get_from_variable(self, var_name: str) -> str:
        """从变量获取数据"""
        var_name = var_name.strip()
        if var_name in self.context_variables:
            return str(self.context_variables[var_name])
        else:
            return f"变量 '{var_name}' 不存在"
    
    def _parse_map_config(self, user_input: str) -> Dict[str, Any]:
        """从用户输入中解析地图配置参数"""
        config = {
            # 绘制范围
            "bounds": {
                "west": 116.0,  # 默认北京范围
                "east": 117.0,
                "south": 39.5,
                "north": 40.5
            },
            # 地图样式
            "style": {
                "background_color": "#f0f8ff",  # 默认浅蓝色背景
                "grid_color": "#cccccc",        # 网格线颜色
                "grid_alpha": 0.3,              # 网格透明度
                "border_color": "#333333",      # 边框颜色
                "border_width": 2               # 边框宽度
            },
            # 地图元素
            "elements": {
                "show_grid": True,              # 显示经纬网
                "show_scale": True,             # 显示比例尺
                "show_north_arrow": True,       # 显示指北针
                "show_title": True,             # 显示标题
                "show_legend": True             # 显示图例
            },
            # 输出设置
            "output": {
                "format": "png",                # 输出格式
                "dpi": 300,                     # 分辨率
                "width": 12,                    # 图片宽度(英寸)
                "height": 8                     # 图片高度(英寸)
            },
            # 标题和描述
            "title": "地理专题图",
            "subtitle": "",
            "description": ""
        }
        
        # 解析绘制范围
        bounds_patterns = [
            r"范围[：:]\s*([0-9.]+)\s*[,，]\s*([0-9.]+)\s*[,，]\s*([0-9.]+)\s*[,，]\s*([0-9.]+)",
            r"边界[：:]\s*([0-9.]+)\s*[,，]\s*([0-9.]+)\s*[,，]\s*([0-9.]+)\s*[,，]\s*([0-9.]+)",
            r"([0-9.]+)\s*[,，]\s*([0-9.]+)\s*[,，]\s*([0-9.]+)\s*[,，]\s*([0-9.]+)\s*的?范围"
        ]
        
        for pattern in bounds_patterns:
            match = re.search(pattern, user_input)
            if match:
                config["bounds"] = {
                    "west": float(match.group(1)),
                    "east": float(match.group(2)),
                    "south": float(match.group(3)),
                    "north": float(match.group(4))
                }
                break
        
        # 解析比例尺
        scale_patterns = [
            r"比例尺[：:]\s*1[：:]([0-9]+)",
            r"1[：:]([0-9]+)\s*比例尺",
            r"比例尺\s*([0-9]+)"
        ]
        
        for pattern in scale_patterns:
            match = re.search(pattern, user_input)
            if match:
                config["elements"]["scale_value"] = int(match.group(1))
                break
        
        # 解析指北针
        if any(keyword in user_input for keyword in ["指北针", "北向", "方向"]):
            config["elements"]["show_north_arrow"] = True
        elif any(keyword in user_input for keyword in ["不显示指北针", "隐藏指北针"]):
            config["elements"]["show_north_arrow"] = False
        
        # 解析经纬网
        if any(keyword in user_input for keyword in ["经纬网", "网格", "坐标网"]):
            config["elements"]["show_grid"] = True
        elif any(keyword in user_input for keyword in ["不显示网格", "隐藏网格", "无网格"]):
            config["elements"]["show_grid"] = False
        
        # 解析底图颜色
        color_patterns = [
            r"底图颜色[：:]\s*(#[0-9a-fA-F]{6})",
            r"背景色[：:]\s*(#[0-9a-fA-F]{6})",
            r"背景[：:]\s*(#[0-9a-fA-F]{6})"
        ]
        
        for pattern in color_patterns:
            match = re.search(pattern, user_input)
            if match:
                config["style"]["background_color"] = match.group(1)
                break
        
        # 解析输出格式
        if any(keyword in user_input for keyword in ["gif", "动画", "动态"]):
            config["output"]["format"] = "gif"
        elif any(keyword in user_input for keyword in ["png", "静态"]):
            config["output"]["format"] = "png"
        
        # 解析标题
        title_patterns = [
            r"标题[：:]\s*([^，。\n]+)",
            r"名为[：:]\s*([^，。\n]+)",
            r"制作[：:]\s*([^，。\n]+)"
        ]
        
        for pattern in title_patterns:
            match = re.search(pattern, user_input)
            if match:
                config["title"] = match.group(1).strip()
                break
        
        return config
    
    def _execute_create_geographic_map(self, config_json: str) -> str:
        """执行地理专题图制作"""
        try:
            if not image_rendering_tools:
                return "错误：地图渲染工具不可用，请检查相关依赖"
            
            # 解析配置参数
            import json
            try:
                config = json.loads(config_json)
            except json.JSONDecodeError:
                return "错误：配置参数格式不正确，请使用有效的JSON格式"
            
            # 根据输出格式选择制作方法
            if config.get("output", {}).get("format", "png") == "gif":
                # 创建动画地图
                map_path = image_rendering_tools.create_animation_gif(
                    duration=config.get("output", {}).get("duration", 2.0)
                )
                format_desc = "动画地图 (GIF格式)"
            else:
                # 创建静态地图
                map_path = image_rendering_tools.create_geo_chart(
                    title=config.get("title", "地理专题图")
                )
                format_desc = "静态地图 (PNG格式)"
            
            # 构建响应信息
            response = f"✅ 已成功制作{format_desc}！\n\n"
            response += f"📁 **文件路径：** {map_path}\n"
            response += f"📏 **绘制范围：** {config.get('bounds', {})}\n"
            response += f"🎨 **底图颜色：** {config.get('style', {}).get('background_color', '#f0f8ff')}\n"
            response += f"📍 **地图元素：**\n"
            
            elements = config.get("elements", {})
            if elements.get("show_grid", True):
                response += "  • ✅ 经纬网\n"
            if elements.get("show_scale", True):
                response += "  • ✅ 比例尺\n"
            if elements.get("show_north_arrow", True):
                response += "  • ✅ 指北针\n"
            if elements.get("show_title", True):
                response += "  • ✅ 标题\n"
            
            response += f"📊 **输出设置：** {config.get('output', {}).get('format', 'png').upper()}格式，{config.get('output', {}).get('dpi', 300)}DPI\n"
            
            return response
            
        except Exception as e:
            return f"制作地理专题图失败: {str(e)}"
    
    def process_request(self, user_input: str) -> str:
        """处理用户请求"""
        try:
            logger.info(f"处理用户请求: {user_input}")
            user_input_lower = user_input.lower()
            
            # 检查是否是地图制作请求
            if any(keyword in user_input_lower for keyword in ['地图', '专题图', '制作', '生成', '创建']):
                # 解析配置参数
                config = self._parse_map_config(user_input)
                
                # 构建配置JSON
                import json
                config_json = json.dumps(config, ensure_ascii=False, indent=2)
                
                # 调用地图制作工具
                return self._execute_create_geographic_map(config_json)
            
            else:
                return """🗺️ 您好！我是地理专题图制作专家，可以帮您制作高质量的地理专题图。

🎯 **支持的功能：**
• 📍 自定义绘制范围（经纬度坐标）
• 📏 可配置比例尺（1:1000 - 1:1000000）
• 🧭 指北针显示控制
• 🎨 底图颜色自定义
• 🌐 经纬网显示控制
• 📊 多种输出格式（PNG/GIF）

💡 **使用示例：**
• "制作北京市地理专题图，范围：116.0,117.0,39.5,40.5，比例尺1:50000"
• "生成上海地图，底图颜色：#e6f3ff，显示指北针和经纬网"
• "创建动画地图，标题：城市发展变化图"

🔧 **配置参数说明：**
• **绘制范围：** 西经,东经,南纬,北纬
• **比例尺：** 1:数字（如1:50000）
• **底图颜色：** 十六进制颜色代码（如#f0f8ff）
• **输出格式：** PNG（静态）或GIF（动画）

请告诉我您想要制作什么样的地理专题图，我会为您配置最合适的参数！"""
            
        except Exception as e:
            logger.error(f"处理请求失败: {e}")
            return f"""❌ 处理请求时出现错误：

**错误详情：** {str(e)}

🔧 **建议操作：**
1. 检查输入格式是否正确
2. 确认配置参数有效
3. 重试您的请求
4. 如问题持续，请联系技术支持

💡 **提示：** 您可以尝试使用更简单的请求来测试系统功能，比如直接输入"制作地图"。"""
    
    def clear_context(self):
        """清空上下文"""
        try:
            self.memory.clear()
            self.context_variables.clear()
            logger.info("上下文已清空")
        except Exception as e:
            logger.error(f"清空上下文失败: {e}")

# 全局agent实例
geographic_map_agent = GeographicMapAgent() 